#!/usr/bin/env python3
"""
进度劫持管理器 - 独立的劫持逻辑管理
"""

from .custom_logger import get_logger

# 为这个模块创建独立的logger
logger = get_logger("progress_hijacker")

class ProgressHijacker:
    """进度劫持管理器 - 管理训练进度的劫持和还原"""
    
    def __init__(self, trainer_instance):
        """初始化劫持管理器
        Args:
            trainer_instance: 训练器实例，用于调用update_progress方法
        """
        self.trainer_instance = trainer_instance
        self.is_hijacked = False
        self.last_reported_step = 0  # 记录上次报告的步数，避免重复报告
        self.step_source = "unknown"  # 记录步数来源：accelerator或tqdm
        self.batch_count = 0  # 记录batch数量，用于处理梯度累积
        self.gradient_accumulation_steps = 1  # 梯度累积步数
    
    def setup_hijacking(self):
        """设置进度劫持"""
        if self.is_hijacked:
            logger.debug("ℹ️  进度劫持已设置，跳过")
            return
        
        # 重置步数状态
        self.last_reported_step = 0
        self.step_source = "unknown"
        self.batch_count = 0
        # 尝试从trainer_instance获取梯度累积步数
        if hasattr(self.trainer_instance, 'parser') and hasattr(self.trainer_instance.parser, 'gradient_accumulation_steps'):
            self.gradient_accumulation_steps = self.trainer_instance.parser.gradient_accumulation_steps
        logger.info(f"🔄 重置步数状态，准备开始新的训练，梯度累积步数: {self.gradient_accumulation_steps}")
        
        try:
            from library import train_util
            
            # 检查是否已经劫持过
            if hasattr(train_util, '_original_prepare_accelerator'):
                logger.debug("ℹ️  train_util.prepare_accelerator函数已被劫持，跳过")
                return
            
            # 检查原始函数是否存在且可调用
            if not hasattr(train_util, 'prepare_accelerator'):
                logger.warning("⚠️  train_util.prepare_accelerator函数不存在，无法劫持")
                return
            
            original_func = getattr(train_util, 'prepare_accelerator', None)
            if not callable(original_func):
                logger.warning("⚠️  train_util.prepare_accelerator不是可调用函数，无法劫持")
                return
            
            # 保存原始函数
            train_util._original_prepare_accelerator = original_func
            
            def patched_prepare_accelerator(args):
                """劫持后的prepare_accelerator函数"""
                # 调用原始函数创建accelerator
                accelerator = train_util._original_prepare_accelerator(args)
                
                logger.info("✅ accelerator创建成功，开始劫持log方法")
                
                # 立即劫持accelerator.log方法
                if hasattr(accelerator, 'log'):
                    # 保存原始log方法
                    accelerator._original_log = accelerator.log
                    
                    def patched_log(logs, step=None, **kwargs):
                        """劫持后的log方法"""
                        # 调用原始log方法
                        result = accelerator._original_log(logs, step, **kwargs)
                        
                        # 添加调试日志
                        logger.debug(f"🔄 劫持的log方法被调用: step={step}, logs={logs}")
                        
                        # 从logs中提取进度信息
                        if step is not None and isinstance(logs, dict):
                            # 统一使用accelerator.log的step参数作为主要步数来源
                            self.trainer_instance.current_step = step
                            self.step_source = "accelerator"
                            
                            # 从logs中提取损失值
                            current_loss = 0.0
                            if "loss/current" in logs:
                                current_loss = logs["loss/current"]
                            elif "loss" in logs:
                                current_loss = logs["loss"]
                            elif "avr_loss" in logs:
                                current_loss = logs["avr_loss"]

                            
                            # 检查并修正total_steps
                            if self.trainer_instance.total_steps <= 0:
                                # 如果total_steps未设置，尝试从logs中获取
                                if "total_steps" in logs:
                                    self.trainer_instance.total_steps = logs["total_steps"]
                                elif "max_train_steps" in logs:
                                    self.trainer_instance.total_steps = logs["max_train_steps"]
                                else:
                                    # 如果都没有，使用默认值
                                    self.trainer_instance.total_steps = 1000
                                    logger.warning(f"⚠️ 无法获取total_steps，使用默认值: {self.trainer_instance.total_steps}")
                            
                            # 计算实际训练步数
                            effective_step = self.calculate_effective_step(step)
                            
                            # 验证步数一致性
                            if self.validate_step_consistency(effective_step, self.trainer_instance.total_steps):
                                # 避免重复报告相同步数
                                if effective_step != self.last_reported_step:
                                    logger.debug(f"📊 提取到进度信息: 报告步数={step}, 实际步数={effective_step}, loss={current_loss}, total_steps={self.trainer_instance.total_steps}, 来源={self.step_source}")
                                    
                                    # 更新进度（使用实际步数）
                                    self.trainer_instance.update_progress(
                                        step=effective_step,
                                        total_steps=self.trainer_instance.total_steps,
                                        loss=current_loss
                                    )
                                    
                                    self.last_reported_step = effective_step
                            else:
                                logger.warning(f"⚠️ 步数验证失败，跳过更新: 报告步数={step}, 实际步数={effective_step}, total_steps={self.trainer_instance.total_steps}")
                            
                            # 计算进度百分比
                            if self.trainer_instance.total_steps > 0:
                                progress_percent = (effective_step / self.trainer_instance.total_steps) * 100
                                logger.info(f"训练进度: {effective_step}/{self.trainer_instance.total_steps} ({progress_percent:.1f}%) - 损失: {current_loss:.4f}")
                            
                            # 检查是否训练完成，如果是则自动还原劫持
                            if effective_step >= self.trainer_instance.total_steps:
                                logger.info("🎉 训练完成！自动还原劫持")
                                self.restore_hijacking()
                            # 检查是否训练异常结束（比如提前停止）
                            elif effective_step > 0 and "loss" in logs and logs.get("loss", 0) == 0:
                                # 如果损失为0且不是第一步，可能是训练异常结束
                                logger.warning("⚠️ 检测到可能的训练异常结束，自动还原劫持")
                                self.restore_hijacking()
                        else:
                            logger.debug(f"⚠️ log方法被调用但step为None或logs不是dict: step={step}, logs={logs}")
                        
                        return result
                    
                    # 替换log方法
                    accelerator.log = patched_log
                    logger.info("✅ 已劫持accelerator.log()方法，开始监控训练进度")
                else:
                    logger.warning("⚠️ accelerator没有log方法，无法劫持")
                
                # 劫持tqdm的progress_bar.set_postfix方法
                try:
                    import tqdm
                    
                    # 保存原始set_postfix方法
                    if not hasattr(tqdm.tqdm, '_original_set_postfix'):
                        tqdm.tqdm._original_set_postfix = tqdm.tqdm.set_postfix
                    
                    # 保存原始update方法
                    if not hasattr(tqdm.tqdm, '_original_update'):
                        tqdm.tqdm._original_update = tqdm.tqdm.update
                    
                    # 保存对当前训练器实例的引用
                    trainer_instance = self.trainer_instance
                    
                    def patched_set_postfix(self, **kwargs):
                        """劫持后的set_postfix方法"""
                        # 调用原始set_postfix方法
                        result = tqdm.tqdm._original_set_postfix(self, **kwargs)
                        
                        # 从kwargs中提取进度信息
                        if kwargs and isinstance(kwargs, dict):
                            # 检查是否是训练相关的进度条（不是缓存阶段）
                            is_training_bar = True
                            if hasattr(self, 'desc') and self.desc:
                                desc = str(self.desc).lower()
                                # 如果是缓存相关的进度条，跳过
                                if any(keyword in desc for keyword in ['caching', 'cache', 'latents', 'text encoder']):
                                    is_training_bar = False
                            
                            if is_training_bar:
                                # 尝试从progress_bar的描述中提取步数信息
                                current_step = 0
                                if hasattr(self, 'n') and self.n is not None:
                                    current_step = self.n
                                
                                # 从kwargs中提取损失值
                                current_loss = 0.0
                                if "avr_loss" in kwargs:
                                    current_loss = kwargs["avr_loss"]
                                elif "loss" in kwargs:
                                    current_loss = kwargs["loss"]
                                
                                # 只有当accelerator.log没有提供step时才使用tqdm的step
                                # 这样可以避免步数冲突
                                if current_step > 0:
                                    trainer_instance.current_step = current_step
                                    trainer_instance.progress_hijacker.step_source = "tqdm"
                                    
                                    # 计算实际训练步数
                                    effective_step = trainer_instance.progress_hijacker.calculate_effective_step(current_step)
                                    
                                    # 验证步数一致性
                                    if trainer_instance.progress_hijacker.validate_step_consistency(effective_step, trainer_instance.total_steps):
                                        # 避免重复报告相同步数
                                        if effective_step != trainer_instance.progress_hijacker.last_reported_step:
                                            logger.debug(f"📊 tqdm提取到进度信息: 报告步数={current_step}, 实际步数={effective_step}, loss={current_loss}, total_steps={trainer_instance.total_steps}, 来源={trainer_instance.progress_hijacker.step_source}")
                                            
                                            trainer_instance.update_progress(
                                                step=effective_step,
                                                total_steps=trainer_instance.total_steps,
                                                loss=current_loss
                                            )
                                            
                                            trainer_instance.progress_hijacker.last_reported_step = effective_step
                                            
                                            # 计算进度百分比
                                            if trainer_instance.total_steps > 0:
                                                progress_percent = (effective_step / trainer_instance.total_steps) * 100
                                                logger.info(f"训练进度: {effective_step}/{trainer_instance.total_steps} ({progress_percent:.1f}%) - 损失: {current_loss:.4f}")
                                            
                                            # 检查是否训练完成
                                            if effective_step >= trainer_instance.total_steps:
                                                logger.info("🎉 训练完成！自动还原劫持")
                                                trainer_instance.progress_hijacker.restore_hijacking()
                                    else:
                                        logger.warning(f"⚠️ tqdm步数验证失败，跳过更新: 报告步数={current_step}, 实际步数={effective_step}, total_steps={trainer_instance.total_steps}")
                        
                        return result
                    
                    def patched_update(self, n=1, *args, **kwargs):
                        """劫持后的update方法，用于监控缓存阶段"""
                        # 调用原始update方法
                        result = tqdm.tqdm._original_update(self, n, *args, **kwargs)
                        
                        # 检查是否有损失值来判断是否为训练进度
                        has_loss = kwargs.get('avr_loss') is not None or kwargs.get('loss') is not None
                        
                        # 有损失值 = 训练进度，没有损失值 = 训练准备进度（缓存等）
                        is_cache_phase = not has_loss
                        
                        if is_cache_phase:
                            current_step = 0
                            if hasattr(self, 'n') and self.n is not None:
                                current_step = self.n
                            total_steps = 0
                            if hasattr(self, 'total') and self.total is not None:
                                total_steps = self.total
                            
                            if current_step > 0 and total_steps > 0:
                                progress_percent = (current_step / total_steps) * 100
                                logger.info(f"缓存进度: {current_step}/{total_steps} ({progress_percent:.1f}%) - {self.desc}")
                                
                                # 推送缓存进度到前端，loss传入-1表示非训练阶段
                                try:
                                    trainer_instance.update_progress(
                                        step=current_step,
                                        total_steps=total_steps,
                                        loss=-1,  # 使用-1表示非训练阶段
                                        status="缓存中"
                                    )
                                except Exception as e:
                                    logger.warning(f"推送缓存进度到前端失败: {e}")
                        
                        return result
                    
                    # 替换方法
                    tqdm.tqdm.set_postfix = patched_set_postfix
                    tqdm.tqdm.update = patched_update
                    logger.info("✅ 已劫持tqdm.set_postfix()和tqdm.update()方法，开始监控训练和缓存进度")
                    
                except Exception as e:
                    logger.error(f"劫持过程中出现异常: {e}")
                
                return accelerator
            
            # 替换prepare_accelerator函数
            train_util.prepare_accelerator = patched_prepare_accelerator
            logger.info("✅ 已劫持train_util.prepare_accelerator函数")
            
            # 标记为已劫持
            self.is_hijacked = True
            
        except Exception as e:
            logger.error(f"[ERROR] 设置进度劫持失败: {e}")
    
    def validate_step_consistency(self, current_step, total_steps):
        """验证步数一致性
        Args:
            current_step: 当前步数
            total_steps: 总步数
        Returns:
            bool: 步数是否合理
        """
        if current_step < 0:
            logger.warning(f"⚠️ 当前步数为负数: {current_step}")
            return False
        
        if total_steps <= 0:
            logger.warning(f"⚠️ 总步数无效: {total_steps}")
            return False
        
        if current_step > total_steps:
            logger.warning(f"⚠️ 当前步数超过总步数: {current_step} > {total_steps}")
            return False
        
        # 检查步数跳跃是否过大（考虑梯度累积）
        if self.last_reported_step > 0:
            step_diff = current_step - self.last_reported_step
            # 如果梯度累积步数大于1，允许更大的跳跃
            max_allowed_jump = max(100, self.gradient_accumulation_steps * 10)
            if step_diff > max_allowed_jump:
                logger.warning(f"⚠️ 步数跳跃过大: {self.last_reported_step} -> {current_step} (跳跃{step_diff}步，梯度累积步数: {self.gradient_accumulation_steps})")
        
        return True
    
    def calculate_effective_step(self, reported_step):
        """计算实际训练步数（考虑梯度累积）
        Args:
            reported_step: 报告的步数
        Returns:
            int: 实际训练步数
        """
        # 根据源码分析，accelerator.log的step是实际的优化步数
        # 而tqdm的step可能是batch计数
        if self.step_source == "accelerator":
            # accelerator.log的step已经是考虑了梯度累积的
            return reported_step
        elif self.step_source == "tqdm":
            # tqdm的step可能是batch计数，需要除以梯度累积步数
            return reported_step // self.gradient_accumulation_steps
        else:
            return reported_step
    
    def restore_hijacking(self):
        """还原进度劫持"""
        if not self.is_hijacked:
            logger.debug("ℹ️  进度劫持未设置，无需还原")
            return
        
        try:
            from library import train_util
            
            # 检查是否已经劫持过
            if hasattr(train_util, '_original_prepare_accelerator'):
                # 检查原始函数是否有效
                original_func = getattr(train_util, '_original_prepare_accelerator', None)
                if original_func is not None and callable(original_func):
                    # 还原原始函数
                    train_util.prepare_accelerator = original_func
                    
                    # 删除原始函数引用
                    delattr(train_util, '_original_prepare_accelerator')
                    
                    logger.info("✅ 已还原train_util.prepare_accelerator函数")
                else:
                    logger.warning("⚠️ 原始函数无效或为空，跳过还原")
            else:
                logger.debug("ℹ️  train_util.prepare_accelerator函数未被劫持，无需还原")
            
            # 还原tqdm的set_postfix方法
            try:
                import tqdm
                if hasattr(tqdm.tqdm, '_original_set_postfix'):
                    tqdm.tqdm.set_postfix = tqdm.tqdm._original_set_postfix
                    delattr(tqdm.tqdm, '_original_set_postfix')
                    logger.info("✅ 已还原tqdm.set_postfix()方法")
                if hasattr(tqdm.tqdm, '_original_update'):
                    tqdm.tqdm.update = tqdm.tqdm._original_update
                    delattr(tqdm.tqdm, '_original_update')
                    logger.info("✅ 已还原tqdm.update()方法")
            except Exception as e:
                logger.warning(f"⚠️ 还原tqdm方法时出错: {e}")
            
            # 标记为未劫持
            self.is_hijacked = False
                
        except Exception as e:
            logger.warning(f"⚠️ 还原进度劫持时出错: {e}") 